package org.zero.ganymed;

import ch.ethz.ssh2.ChannelCondition;
import ch.ethz.ssh2.Connection;
import ch.ethz.ssh2.Session;
import ch.ethz.ssh2.StreamGobbler;
import java.io.IOException;
import java.io.InputStream;
import java.io.PrintWriter;
import java.nio.charset.Charset;
import lombok.extern.slf4j.Slf4j;

/**
 * @author v_wangyufa (cnzeropro@163.com)
 * @date 2022/2/17 11:28
 */
@Slf4j
public class RemoteShellExecutor {

    private Connection connection;

    private final String ip;

    private final String user;

    private final String password;

    private final String charset = Charset.defaultCharset().toString();

    private static final int TIME_OUT = 1000 * 60 * 5;

    public RemoteShellExecutor(String ip, String user, String password) {
        this.ip = ip;
        this.user = user;
        this.password = password;

        init();
    }

    /**
     * 初始化
     *
     * @throws IOException
     * @resulturn
     */
    private void init() {
        log.info("正在初始化连接...");
        connection = new Connection(ip);
    }


    /**
     * 登录认证
     *
     * @throws IOException
     * @resulturn
     */
    private boolean login() throws IOException {
        connection.connect();
        return connection.authenticateWithPassword(user, password);
    }

    /**
     * 执行脚本
     *
     * @param cmds
     * @throws Exception
     * @resulturn
     */
    public int exec(String cmds) throws Exception {
        InputStream stdOut = null;
        InputStream stdErr = null;
        String outStd = "";
        String outErr = "";
        Integer result = -1;
        try {
            if (!connection.isAuthenticationComplete() && !login()) {
                String errStr = "登录远程机器失败：" + ip;
                log.error(errStr);
                throw new RuntimeException(errStr);
            }

            log.info("登录验证成功...");

            log.info("开启 SSH 会话...");
            Session session = connection.openSession();
            log.info("开始执行命令：{}", cmds);

            // 用方法execCommand执行Shell命令的时候，有时会遇到获取不全环境变量
            // session.execCommand(cmds);

            // 所以换另外一种方式来执行Shell命令：建立虚拟终端
            session.requestPTY("bash");
            session.startShell();
            PrintWriter in = new PrintWriter(session.getStdin());
            in.println(cmds);
            in.println("exit");
            in.close();

            // 等待命令执行
            session.waitForCondition(ChannelCondition.CLOSED | ChannelCondition.EOF | ChannelCondition.EXIT_STATUS, TIME_OUT);

            log.info("正在处理标准输出流...");
            stdOut = new StreamGobbler(session.getStdout());
            outStd = processStream(stdOut, charset);
            log.info("outStr={}", outStd);

            log.info("正在处理标准错误流...");
            stdErr = new StreamGobbler(session.getStderr());
            outErr = processStream(stdErr, charset);
            log.info("outStr={}", outErr);

            result = session.getExitStatus();
        } finally {
            if (stdOut != null) {
                stdOut.close();
            }
            if (stdErr != null) {
                stdErr.close();
            }
        }

        return result;
    }

    /**
     * 处理信息流
     *
     * @param in
     * @param charset
     * @return
     * @throws Exception
     */
    private String processStream(InputStream in, String charset) throws Exception {
        byte[] buf = new byte[1024];
        StringBuilder sb = new StringBuilder();
        while (in.read(buf) != -1) {
            sb.append(new String(buf, charset));
        }
        return sb.toString();
    }

    /**
     * 关闭 SSH 连接
     */
    public void closeConnection() {
        if (connection != null) {
            connection.close();
        }
    }
}
